import sklearn_crfsuite
from sklearn_crfsuite import scorers
from sklearn_crfsuite import metrics

# 用于生成特征的函数
def word2features(sent, i):
    word = sent[i][0]
    features = {
        'bias': 1.0,                    # 特征实际上是一个常数偏置项，用于确保模型不会预测所有0。它可以被看作是线性模型方程中的截距项
        'word.lower()': word.lower(),   # 相同单词的不同大小写形式应具有相同或相似的标签
        'word[-3:]': word[-3:],         # 通常用于捕获如词缀或词根等语言特性
    }
    return features

# 生成句子的特征
def sent2features(sent):
    return [word2features(sent, i) for i in range(len(sent))]

# 生成句子的标签
def sent2labels(sent):
    return [label for token, label in sent]

# 示例数据（单词，标签）
train_sents = [
[('dog', 'NN'), ('chases', 'VB'), ('cat', 'NN')],
[('cat', 'NN'), ('climbs', 'VB'), ('tree', 'NN')]
]

test_sents = [
[('dog', 'NN'), ('catches', 'VB'), ('ball', 'NN')],
[('cat', 'NN'), ('sits', 'VB'), ('here', 'RB')],
[('cat', 'NN'), ('is', 'AA'), ('here', 'animal')],
]

# 生成特征和标签
X_train = [sent2features(s) for s in train_sents]
y_train = [sent2labels(s) for s in train_sents]

X_test = [sent2features(s) for s in test_sents]
y_test = [sent2labels(s) for s in test_sents]

# 创建并训练CRF模型
crf = sklearn_crfsuite.CRF(
    algorithm='lbfgs',
    c1=0.1,
    c2=0.1,
    max_iterations=100,
    all_possible_transitions=True
)
crf.fit(X_train, y_train)

# 在测试集上进行预测
y_pred = crf.predict(X_test)

# 计算评估指标
labels = list(crf.classes_)
print(labels)
metrics.flat_f1_score(y_test, y_pred, average='weighted', labels=labels)

# 输出预测结果
print(y_pred)